from __future__ import print_function

import argparse
import bz2
import gzip
import io
import lzma
import os
import sys
import zlib
from pathlib import Path

CHUNK_SIZE = 1024 * 1024  # 1 MiB

def _ensure_level(level):
    try:
        lvl = int(level)
    except Exception:
        lvl = 6
    return max(0, min(9, lvl))

def detect_algorithm(header_bytes):
    """根据文件头字节检测压缩算法。

    返回值："gzip" | "bz2" | "lzma" | "zlib"
    未识别时抛出 ValueError。
    """
    hb = header_bytes
    if len(hb) >= 2 and hb[0] == 0x1F and hb[1] == 0x8B:
        return "gzip"
    if hb.startswith(b"BZh"):
        return "bz2"
    if hb.startswith(b"\xFD7zXZ\x00"):
        return "lzma"  # XZ 容器
    if len(hb) >= 2 and hb[0] == 0x78 and ((hb[0] * 256 + hb[1]) % 31) == 0:
        return "zlib"
    raise ValueError("无法自动识别压缩算法（非 gzip/bz2/xz/zlib 格式）")

def _default_extension_for_algo(algo):
    return {
        "gzip": ".gz",
        "bz2": ".bz2",
        "lzma": ".xz",
        "zlib": ".zlib",
    }.get(algo, ".bin")

def encode_file(input_path, output_path=None, algorithm="lzma", level=6):
    """对文件进行熵编码（压缩）。

    algorithm 可选："zlib" | "gzip" | "bz2" | "lzma"
    level 范围：0-9
    返回：输出文件的 Path
    """
    algo = algorithm.lower()
    lvl = _ensure_level(level)
    in_path = Path(input_path)
    if output_path is None:
        ext = _default_extension_for_algo(algo)
        out_path = in_path.with_suffix(in_path.suffix + ext)
    else:
        out_path = Path(output_path)

    if algo not in {"zlib", "gzip", "bz2", "lzma"}:
        raise ValueError("不支持的算法：" + algo)

    out_path.parent.mkdir(parents=True, exist_ok=True)

    if algo == "gzip":
        with open(str(in_path), "rb") as f_in, gzip.open(str(out_path), "wb", compresslevel=lvl) as f_out:
            for chunk in iter(lambda: f_in.read(CHUNK_SIZE), b""):
                f_out.write(chunk)
        return out_path

    if algo == "bz2":
        with open(str(in_path), "rb") as f_in, bz2.open(str(out_path), "wb", compresslevel=lvl) as f_out:
            for chunk in iter(lambda: f_in.read(CHUNK_SIZE), b""):
                f_out.write(chunk)
        return out_path

    if algo == "lzma":
        # 使用 XZ 容器，利于自动识别
        with open(str(in_path), "rb") as f_in, lzma.open(str(out_path), "wb", format=lzma.FORMAT_XZ, preset=lvl) as f_out:
            for chunk in iter(lambda: f_in.read(CHUNK_SIZE), b""):
                f_out.write(chunk)
        return out_path

    # zlib：流式压缩（生成带 zlib 头的 DEFLATE 数据）
    comp = zlib.compressobj(lvl)
    with open(str(in_path), "rb") as f_in, open(str(out_path), "wb") as f_out:
        for chunk in iter(lambda: f_in.read(CHUNK_SIZE), b""):
            f_out.write(comp.compress(chunk))
        f_out.write(comp.flush())
    return out_path

def decode_file(input_path, output_path=None, algorithm="auto"):
    """对文件进行熵解码（解压）。

    algorithm 可选："auto" | "zlib" | "gzip" | "bz2" | "lzma"
    返回：输出文件的 Path
    """
    algo = algorithm.lower()
    in_path = Path(input_path)
    if output_path is None:
        # 去掉已知扩展；若无法识别则添加 .dec 结尾
        suffixes = in_path.suffixes
        out_path = Path(in_path)
        if suffixes:
            # 移除最后一个后缀
            out_path = in_path.with_suffix("")
            if len(suffixes) >= 2:
                # 保留前面的后缀
                for s in suffixes[:-1]:
                    out_path = out_path.with_suffix(s)
        else:
            out_path = in_path.with_suffix(in_path.suffix + ".dec")
    else:
        out_path = Path(output_path)

    # 自动识别
    if algo == "auto":
        with open(str(in_path), "rb") as f_in:
            header = f_in.read(10)
        try:
            algo = detect_algorithm(header)
        except ValueError as e:
            raise ValueError("auto 识别失败：" + str(e))

    if algo not in {"zlib", "gzip", "bz2", "lzma"}:
        raise ValueError("不支持的算法：" + algo)

    out_path.parent.mkdir(parents=True, exist_ok=True)

    if algo == "gzip":
        with gzip.open(str(in_path), "rb") as f_in, open(str(out_path), "wb") as f_out:
            for chunk in iter(lambda: f_in.read(CHUNK_SIZE), b""):
                f_out.write(chunk)
        return out_path

    if algo == "bz2":
        with bz2.open(str(in_path), "rb") as f_in, open(str(out_path), "wb") as f_out:
            for chunk in iter(lambda: f_in.read(CHUNK_SIZE), b""):
                f_out.write(chunk)
        return out_path

    if algo == "lzma":
        with lzma.open(str(in_path), "rb") as f_in, open(str(out_path), "wb") as f_out:
            for chunk in iter(lambda: f_in.read(CHUNK_SIZE), b""):
                f_out.write(chunk)
        return out_path

    # zlib：流式解压
    decomp = zlib.decompressobj()
    with open(str(in_path), "rb") as f_in, open(str(out_path), "wb") as f_out:
        for chunk in iter(lambda: f_in.read(CHUNK_SIZE), b""):
            f_out.write(decomp.decompress(chunk))
        f_out.write(decomp.flush())
    return out_path

def _parse_args(argv):
    p = argparse.ArgumentParser(description="文件熵编码/解码工具（zlib/gzip/bz2/lzma）")
    sub = p.add_subparsers(dest="cmd")

    p_enc = sub.add_parser("encode", help="对文件进行熵编码（压缩）")
    p_enc.add_argument("input", help="输入文件路径")
    p_enc.add_argument("-o", "--output", help="输出文件路径（默认根据算法追加扩展名）")
    p_enc.add_argument("-a", "--algo", default="lzma", choices=["zlib", "gzip", "bz2", "lzma"], help="压缩算法")
    p_enc.add_argument("-l", "--level", default=6, type=int, help="压缩等级 0-9（默认 6）")

    p_dec = sub.add_parser("decode", help="对文件进行熵解码（解压）")
    p_dec.add_argument("input", help="输入文件路径")
    p_dec.add_argument("-o", "--output", help="输出文件路径（默认移除已知扩展）")
    p_dec.add_argument("-a", "--algo", default="auto", choices=["auto", "zlib", "gzip", "bz2", "lzma"], help="解压算法")

    return p.parse_args(argv)

def main(argv=None):
    ns = _parse_args(argv or sys.argv[1:])
    try:
        if ns.cmd == "encode":
            outp = encode_file(ns.input, ns.output, ns.algo, ns.level)
            print("编码完成：" + str(outp))
        elif ns.cmd == "decode":
            outp = decode_file(ns.input, ns.output, ns.algo)
            print("解码完成：" + str(outp))
        else:
            raise RuntimeError("未知命令")
        return 0
    except Exception as e:
        print("错误：" + str(e), file=sys.stderr)
        return 1

if __name__ == "__main__":
    sys.exit(main())